#!/usr/bin/env python3
import os
import sys
import time
import argparse

class HardwarePWM:
    def __init__(self, chip, channel):
        self.chip = chip
        self.channel = channel
        self.pwm_path = f"/sys/class/pwm/{chip}/pwm{channel}"
        
    def export_channel(self):
        """导出PWM通道"""
        if not os.path.exists(self.pwm_path):
            try:
                with open(f"/sys/class/pwm/{self.chip}/export", "w") as f:
                    f.write(str(self.channel))
                print(f"已导出 {self.chip} pwm{self.channel}")
                time.sleep(0.1)  # 等待设备创建
            except Exception as e:
                print(f"导出失败: {e}")
                return False
        return True
    
    def unexport_channel(self):
        """取消导出PWM通道"""
        if os.path.exists(self.pwm_path):
            try:
                with open(f"/sys/class/pwm/{self.chip}/unexport", "w") as f:
                    f.write(str(self.channel))
                print(f"已取消导出 {self.chip} pwm{self.channel}")
            except Exception as e:
                print(f"取消导出失败: {e}")
    
    def set_period(self, period_ns):
        """设置周期（纳秒）"""
        try:
            with open(f"{self.pwm_path}/period", "w") as f:
                f.write(str(period_ns))
        except Exception as e:
            print(f"设置周期失败: {e}")
    
    def set_duty_cycle(self, duty_cycle_ns):
        """设置占空比（纳秒）"""
        try:
            with open(f"{self.pwm_path}/duty_cycle", "w") as f:
                f.write(str(duty_cycle_ns))
        except Exception as e:
            print(f"设置占空比失败: {e}")
    
    def enable(self):
        """启用PWM"""
        try:
            with open(f"{self.pwm_path}/enable", "w") as f:
                f.write("1")
        except Exception as e:
            print(f"启用PWM失败: {e}")
    
    def disable(self):
        """禁用PWM"""
        try:
            with open(f"{self.pwm_path}/enable", "w") as f:
                f.write("0")
        except Exception as e:
            print(f"禁用PWM失败: {e}")
    
    def set_frequency_duty(self, frequency_hz, duty_percent):
        """设置频率和占空比百分比"""
        period_ns = int(1e9 / frequency_hz)  # 周期(纳秒)
        duty_cycle_ns = int(period_ns * duty_percent / 100)  # 占空比(纳秒)
        
        self.set_period(period_ns)
        self.set_duty_cycle(duty_cycle_ns)
        print(f"频率: {frequency_hz}Hz, 占空比: {duty_percent}%")

def list_pwm_chips():
    """列出所有PWM芯片"""
    print("可用的PWM芯片:")
    pwm_chips = [d for d in os.listdir("/sys/class/pwm") if d.startswith("pwmchip")]
    for chip in pwm_chips:
        chip_path = f"/sys/class/pwm/{chip}"
        try:
            with open(f"{chip_path}/npwm", "r") as f:
                npwm = f.read().strip()
            print(f"  {chip}: {npwm} 个通道")
        except:
            print(f"  {chip}: 无法读取信息")

def main():
    # 创建解析器
    parser = argparse.ArgumentParser(
        description='硬件PWM控制脚本',
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    
    # 必需参数
    parser.add_argument('chip', nargs='?', help='PWM芯片名称，如 pwmchip0')
    
    # 可选参数（短参数和长参数都支持）
    parser.add_argument('-c', '--channel', type=int, default=0, help='PWM通道 (默认: 0)')
    parser.add_argument('-f', '--frequency', type=float, default=1000, help='频率 (Hz) (默认: 1000)')
    parser.add_argument('-d', '--duty', type=float, default=50, help='占空比百分比 (默认: 50)')
    parser.add_argument('-t', '--time', type=float, default=5, help='运行时间秒 (默认: 5)')
    parser.add_argument('-l', '--list', action='store_true', help='列出所有PWM芯片')
    
    # 解析参数
    args = parser.parse_args()
    
    if args.list:
        list_pwm_chips()
        return
    
    if not args.chip:
        print("错误: 必须指定PWM芯片!")
        print("\n使用示例:")
        print("  python3 pwm_control.py pwmchip0                    # 使用默认参数")
        print("  python3 pwm_control.py pwmchip0 -c 1 -f 2000 -d 75 -t 10")
        print("  python3 pwm_control.py pwmchip1 -f 500 -d 25       # 500Hz, 25%占空比")
        print("  python3 pwm_control.py -l                          # 列出所有PWM芯片")
        return
    
    # 创建PWM对象
    pwm = HardwarePWM(args.chip, args.channel)
    
    try:
        # 导出通道
        if not pwm.export_channel():
            return
        
        # 设置PWM参数
        print(f"设置 {args.chip} 通道 {args.channel}:")
        pwm.set_frequency_duty(args.frequency, args.duty)
        pwm.enable()
        
        print(f"PWM运行中... 持续 {args.time} 秒")
        print("按 Ctrl+C 停止")
        time.sleep(args.time)
        
    except KeyboardInterrupt:
        print("\n用户中断")
    finally:
        # 清理
        pwm.disable()
        pwm.unexport_channel()

if __name__ == "__main__":
    main()